/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.analysys.presto.connector.hbase;

import io.airlift.log.Logger;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.CellUtil;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.client.ResultScanner;
import org.apache.hadoop.hbase.client.Table;
import org.apache.hadoop.hbase.util.Bytes;

import java.net.InetAddress;
import java.util.*;
import static java.util.Objects.requireNonNull;

/**
 * Use record.rawCells() api to loop column value, this is 20% faster than result.getValue
 *
 * @author wupeng
 * @date 2019/01/29
 */
public class HBaseScanRecordCursor extends HBaseRecordCursor {

    public static final Logger log = Logger.get(HBaseGetRecordCursor.class);



    /**
     * hbase connection
     */
    private Connection connection;


    /**
     * hbase table
     */
    private Table table;

    private ResultScanner resultScanner;


    private Iterator iterator;


    /**
     * 多行
     */
    private final LinkedList<Object[]> rows;


    private Result currentRecord;

    HBaseScanRecordCursor(List<HBaseColumnHandle> columnHandles,
                          HBaseSplit hbaseSplit,
                          ResultScanner scanner,
                          Map<Integer, HBaseColumnHandle> fieldIndexMap,
                          Table table,
                          Connection connection) {
        this.startTime = System.currentTimeMillis();
        this.rows = new LinkedList<>();
        this.columnHandles = columnHandles;
        this.fieldIndexMap = fieldIndexMap;

        this.rowKeyColName = requireNonNull(hbaseSplit.getRowKeyName(),
                "RowKeyName cannot be null if you want to query by RowKey");

        this.split = hbaseSplit;
        this.table = table;
        this.connection = connection;
        try {
            this.resultScanner = scanner;
            if (resultScanner != null) {
                this.iterator = resultScanner.iterator();
            }
        } catch (Exception ex) {
            log.error(ex, ex.getMessage());
            this.close();
        }
        this.totalBytes = 0L;
    }

    @Override
    public boolean advanceNextPosition() {
        try {
            Result record = (currentRecord != null ? currentRecord : this.getNextRecord());
            if (record == null) {
                InetAddress localhost = InetAddress.getLocalHost();
                log.info("SCAN RECORD. advanceNextPosition: tableName=" + split.getTableName() + ", startRow=" + split.getStartRow()
                        + ", endRow=" + split.getEndRow() + ". READ_DATA_TIME="
                        + (System.currentTimeMillis() - startTime) + " mill secs. recordCount=" + recordCount
                        + ", startTime=" + new Date(startTime).toString() + ", localhost=" + localhost.getHostAddress()
                        + ", specified worker ip: "
                        + (split.getAddresses().size() > 0 ? split.getAddresses().get(0).toString() : ""));
                return false;
            } else {
                Object[] fields = new Object[this.columnHandles.size()];
                ordinalPositionAndFieldsIndexMap.clear();
                int fieldIndex = 0;
                if(split.isExplodeMode()) {
                    this.currentRecord = record;
                    // 需要行转列的模式
                    if(rows.isEmpty()) {
                        // 需要行转列的模式
                        long start = System.currentTimeMillis();
                        // 如果没有行了，则扫描下一个
                        List<String> cols = new ArrayList<>();
                        Map<String, byte[]> resultMap = new HashMap<>();
                        for (Cell cell : record.rawCells()) {
                            if (Bytes.toString(CellUtil.cloneQualifier(cell)).toLowerCase(Locale.ENGLISH).startsWith(split.getExplodeColumn() + "_")) {
                                cols.add(Bytes.toString(CellUtil.cloneValue(cell)));
                            }
                            String colName = Bytes.toString(CellUtil.cloneFamily(cell)) + ":" + Bytes.toString(CellUtil.cloneQualifier(cell));

                            // 字段名称转为小写
                            colName = colName.toLowerCase(Locale.ENGLISH);
                            resultMap.put(colName, CellUtil.cloneValue(cell));
                        }

                        // columnHandles will be empty if sql is like count(*)
                        if (columnHandles.isEmpty()) {
                            if (rows.isEmpty()) {
                                // 行转列完成，下一行开始
                                this.currentRecord = null;
                            }
                            return true;
                        }

                        // 从 record 读取所有的 行
                        for (String col : cols) {
                            Object[] objs = new Object[columnHandles.size()];
                            int i = 0;
                            for (HBaseColumnHandle column : columnHandles) {
                                if (column.isIsRowKey()) {
                                    continue;
                                }
                                String colN = column.getFamily() + ":" + column.getColumnName() + "_" + col;
                                // 上面字段名称是转为小写的，这里要转为小写去获取？
                                byte[] r = resultMap.get(StringUtils.lowerCase(colN));
                                if (r == null) {
                                    // 如果 explode column 拼接字段没有获取到，则采用直接的字段拼接来获取
                                    r = resultMap.get(StringUtils.lowerCase(column.getFamily() + ":" + column.getColumnName()));
                                }
                                Object value = null;
                                if(r == null){
                                    value = null;
                                } else if ("BINARY".equalsIgnoreCase(split.getStorageType())) {
                                    this.totalBytes += r.length;
                                    value = matchValueBinary(column.getColumnType(), r);
                                } else {
                                    this.totalBytes += r.length;
                                    value = matchValueText(column.getColumnType(), r);
                                }

                                objs[i] = value;
                                i++;
                            }
                            // 添加一条数据
                            rows.add(objs);
                        }
                        log.info("read hbase explode row has " + rows.size() + " rows, cost time: " + (System.currentTimeMillis() - start) + " ms.");
                    }

                    // 还没有数据
                    if(rows.isEmpty()) {
                        currentRecord = null;
                        return false;
                    }
                    // 弹出一行数据
                    Object[] row = rows.poll();
                    int i = 0;
                    for (HBaseColumnHandle column : columnHandles) {
                        if(column.isIsRowKey()) {
                            continue;
                        }
                        fields[fieldIndex] = row[i];
                        ordinalPositionAndFieldsIndexMap.put(column.getOrdinalPosition(), fieldIndex);
                        fieldIndex++;
                        i++;
                    }

                    // Handle the value of rowKey column
                    setRowKeyValue2FieldsAry(record, fieldIndex, fields);

                    // 一行读取完成
                    setFieldValues(fields);
                    if(rows.isEmpty()) {
                        // 行转列完成，下一行开始
                        this.currentRecord = null;
                    }
                } else {
                    // columnHandles will be empty if sql is like count(*)
                    if (columnHandles.isEmpty()) {
                        return true;
                    }
                    for (Cell cell : record.rawCells()) {
                        String colName = Bytes.toString(CellUtil.cloneQualifier(cell)).toLowerCase(Locale.ENGLISH);
                        // Fetch the index and type of column by column name
                        HBaseColumnHandle hch = fieldIndexMap.get(colName.hashCode());
                        if (hch == null) {
                            continue;
                        }

                        this.totalBytes += cell.getValueLength();
                        // Set value to fields array
                        if ("BINARY".equalsIgnoreCase(split.getStorageType())) {
                            fields[fieldIndex] = matchValueBinary(hch.getColumnType(), CellUtil.cloneValue(cell));
                        } else {
                            fields[fieldIndex] = matchValueText(hch.getColumnType(), CellUtil.cloneValue(cell));
                        }
                        // Store values in an array for queries
                        ordinalPositionAndFieldsIndexMap.put(hch.getOrdinalPosition(), fieldIndex);
                        fieldIndex++;
                    }
                    // Handle the value of rowKey column
                    setRowKeyValue2FieldsAry(record, fieldIndex, fields);
                    setFieldValues(fields);
                    currentRecord = null;
                }
                return true;
            }
        } catch (Exception ex) {
            log.error(ex, ex.getMessage());
            this.close();
            return false;
        }
    }

    private Result getNextRecord() {
        try {
            if (iterator == null || !iterator.hasNext()) {
                return null;
            } else {
                recordCount++;
                do {
                    Result r = (Result) iterator.next();
                    if (r != null && !r.isEmpty()) {
                        return r;
                    }
                } while (iterator.hasNext());
            }
        } catch (Exception ex) {
            log.error(ex, ex.getMessage());
            this.close();
        }
        return null;
    }

    @Override
    public void close() {
        if (this.resultScanner != null) {
            try {
                this.resultScanner.close();
            } catch (Exception e) {
                log.warn(e.getMessage(), e);
            }
        }

        if(table != null) {
            try {
                table.close();
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }

        if (connection != null) {
            try {
                connection.close();
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
    }

    @Override
    public long getCompletedBytes() {
        return this.totalBytes;
    }

    @Override
    public long getReadTimeNanos() {
        return System.currentTimeMillis() - startTime;
    }

}
